kthena-router: order boosted requests by session completion time#1334
Open
YaoZengzeng wants to merge 7 commits into
Open
kthena-router: order boosted requests by session completion time#1334YaoZengzeng wants to merge 7 commits into
YaoZengzeng wants to merge 7 commits into
Conversation
Contributor
|
[APPROVALNOTIFIER] This PR is NOT APPROVED This pull-request has been approved by: The full list of commands accepted by this bot can be found here. DetailsNeeds approval from an approver in each of these files:Approvers can indicate their approval by writing |
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
Contributor
There was a problem hiding this comment.
Code Review
This pull request updates the session boost queue scheduling strategy to prioritize the most recently arrived boosted requests (LIFO) rather than serving them in FIFO order, maximizing prefix cache hits. Additionally, it removes the fast path that skipped the grace period when a boosted request was already at the head of the queue, ensuring the queue always waits for the full grace window for potential newer same-session follow-ups. These changes are reflected in the documentation, core queue logic, and unit tests. I have no feedback to provide as there are no review comments.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
Contributor
There was a problem hiding this comment.
Pull request overview
Updates the kthena-router per-model queue’s session-boost behavior to prioritize the most recently arrived boosted request (LIFO within the boosted group), improving the likelihood of reusing warm prefix (KV) cache for same-session follow-ups. It also adjusts the grace-period behavior and documentation to match the intended scheduling semantics.
Changes:
- Change session-boost queue ordering so boosted requests are newest-first (LIFO) while non-boosted remain FIFO.
- Remove the “head already boosted” grace fast path so the full grace window is always honored (when enabled).
- Update unit tests and docs to reflect the new ordering and grace behavior.
Reviewed changes
Copilot reviewed 5 out of 5 changed files in this pull request and generated 1 comment.
Show a summary per file
| File | Description |
|---|---|
| pkg/kthena-router/datastore/session_boost_queue.go | Removes the grace fast path and updates grace behavior documentation in code comments. |
| pkg/kthena-router/datastore/session_boost_queue_test.go | Updates assertions to expect newest-first ordering among boosted requests. |
| pkg/kthena-router/datastore/fairness_queue.go | Adjusts heap ordering in session-boost mode to prefer newer boosted requests. |
| docs/proposal/session-boost-strategy.md | Updates proposal to describe newest-first boosted ordering and no grace fast path. |
| docs/kthena/docs/user-guide/session-boost.md | Updates user guide to describe newest-first boosted ordering and full grace behavior. |
💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.
Comment on lines
+179
to
+181
| 1. A request finishes → `Release()` runs → `inflightCount` is decremented → a signal is sent on `releaseCh`. | ||
| 2. The freed slot would normally be claimed immediately by the current heap head (which may be an unrelated, non-boosted request). The grace wait instead **holds that just-freed slot for up to `SessionBoostGracePeriod`**, betting that a same-session follow-up is about to arrive and reuse the warm prefix cache. | ||
| 3. When the wait resolves, `tryBackpressureDequeue` runs and re-checks **both** capacity gates before admitting anyone. A boosted follow-up that arrived during grace is admitted first because boosted requests outrank others in the heap, and only when `inflight < MaxConcurrent` **and** at least one backend pod reports capacity. | ||
| 3. When the wait resolves, `tryBackpressureDequeue` runs and re-checks **both** capacity gates before admitting anyone. A boosted follow-up that arrived during grace is admitted first because boosted requests outrank others in the heap and, among boosted requests, the most-recently-arrived one is at the head — and only when `inflight < MaxConcurrent` **and** at least one backend pod reports capacity. |
YaoZengzeng
changed the title
Comment on lines
+167
to
+171
| │ │ first and, among boosted, the │ | ||
| │ │ newest-arrived wins, so a │ | ||
| │ │ same-session follow-up that │ | ||
| │ │ arrived during the window is │ | ||
| │ │ dequeued first automatically. │ |
Comment on lines
+247
to
+249
| q.MarkSessionRequestCompleted("conv-A") | ||
| q.MarkSessionRequestCompleted("conv-B") | ||
| q.MarkSessionRequestCompleted("conv-C") |
Member
Author
|
@hzxuzhonghu ptal |
hzxuzhonghu
reviewed
Jul 10, 2026
| // SessionCompletedAt is the time the session's previous turn completed, captured | ||
| // when the request is boosted. Boosted requests are ordered by this timestamp | ||
| // (most recent first) so the session with the warmest prefix cache runs first. | ||
| SessionCompletedAt time.Time |
Member
There was a problem hiding this comment.
The name is a little confusing.
Member
Author
There was a problem hiding this comment.
renamed to LastTurnCompletedAt
Comment on lines
+247
to
+249
| q.MarkSessionRequestCompleted("conv-A") | ||
| q.MarkSessionRequestCompleted("conv-B") | ||
| q.MarkSessionRequestCompleted("conv-C") |
Comment on lines
+220
to
224
| // Among boosted: the session that completed most recently wins. conv-B was | ||
| // marked completed after conv-A, so boost-B (u3) is dequeued before boost-A (u2). | ||
| if first.UserID != "u3" { | ||
| t.Errorf("Expected boost-B first (most recently completed session), got %s", first.UserID) | ||
| } |
Comment on lines
+246
to
+249
| // Completion order: conv-A, then conv-B, then conv-C (conv-C most recent). | ||
| q.MarkSessionRequestCompleted("conv-A") | ||
| q.MarkSessionRequestCompleted("conv-B") | ||
| q.MarkSessionRequestCompleted("conv-C") |
Comment on lines
+167
to
+171
| │ │ first and, among boosted, the │ | ||
| │ │ newest-arrived wins, so a │ | ||
| │ │ same-session follow-up that │ | ||
| │ │ arrived during the window is │ | ||
| │ │ dequeued first automatically. │ |
Comment on lines
+187
to
+189
| | Outcome | Trigger | What happens next | | ||
| | ----------------- | ---------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | ||
| | **Grace expires** | The timer fires (`timer.C`). | Stop holding the slot and fall through to the capacity gates, which admit the heap head. If a same-session follow-up arrived during the wait it is now boosted and sits at the head (session with the most recent completion first), so it wins automatically (subject to inflight + backend capacity). | |
Comment on lines
+166
to
+169
| │ │ head. Boosted requests rank │ | ||
| │ │ first and, among boosted, the │ | ||
| │ │ newest-arrived wins, so a │ | ||
| │ │ same-session follow-up that │ |
| | -------------------------- | --------------------------------------------------------------------------------------------- | ------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- | | ||
| | **Skip grace (fast path)** | The heap head is *already* session-boosted when the release fires (`isHeadSessionBoosted()`). | No wait at all — go straight to the capacity gates and admit if they pass. There is nothing to wait for. | | ||
| | **Grace expires** | The timer fires (`timer.C`). | Stop holding the slot and fall through to the capacity gates, which admit the heap head. If a same-session follow-up arrived during the wait it is now boosted and sits at the head, so it wins automatically (subject to inflight + backend capacity). | | ||
| > There is no "head already boosted" fast path. The queue holds the freed slot for the full grace window even when a boosted request is already at the head, because a still-newer same-session follow-up may yet arrive during the window and should be the one to ride the warm prefix cache. |
Comment on lines
+292
to
+294
| - If a boosted (same-session) request arrives during the grace window, it is admitted first when the window ends, because boosted requests outrank others in the queue—and among boosted requests the one whose session completed most recently is at the head. | ||
| - If no boosted request arrives before the window expires, the next non-boosted request proceeds as normal. | ||
| - If the head of the queue is already a boosted request, the grace period is skipped entirely. | ||
| - The queue always waits for the full grace window even when a boosted request is already at the head, because a still-newer same-session follow-up may yet arrive and should be the one to ride the warm prefix cache. |
| │ │ keys: sessionID │ │ (after response sent) │ │ | ||
| │ │ cap: 4096 default│ └─────────────────────────┘ │ | ||
| │ │ key: sessionID │ │ (after response sent) │ │ | ||
| │ │ value: completeAt│ └─────────────────────────┘ │ |
| │ │ grace period, then dequeue the │ | ||
| │ │ head. Boosted requests rank │ | ||
| │ │ first and, among boosted, the │ | ||
| │ │ newest-arrived wins, so a │ |
hzxuzhonghu
reviewed
Jul 10, 2026
| } | ||
| if pq.heap[i].SessionBoost { | ||
| if !pq.heap[i].LastTurnCompletedAt.Equal(pq.heap[j].LastTurnCompletedAt) { | ||
| return pq.heap[i].LastTurnCompletedAt.After(pq.heap[j].LastTurnCompletedAt) |
Member
There was a problem hiding this comment.
Newest-first ordering can starve an older boosted request. Each later-completing session gets a newer timestamp and stays ahead of it, while the queued request keeps the timestamp captured at enqueue, so sustained boosted traffic can make the older request hit SESSION_BOOST_TIMEOUT instead of ever running. Please add bounded overtaking or aging, with a test that inserts a newer completion before each dequeue.
Comment on lines
+387
to
+402
| mutated := false | ||
| for _, q := range pq.heap { | ||
| if q.SessionBoost && !q.boostAged && q.LastTurnCompletedAt.Before(r.LastTurnCompletedAt) { | ||
| q.overtakenCount++ | ||
| if q.overtakenCount >= maxOvertakes { | ||
| q.boostAged = true | ||
| } | ||
| mutated = true | ||
| } | ||
| } | ||
| if mutated { | ||
| pq.heap = append(pq.heap, r) | ||
| heap.Init(pq) | ||
| } else { | ||
| heap.Push(pq, r) | ||
| } |
| │ │ grace period, then dequeue the │ | ||
| │ │ head. Boosted requests rank │ | ||
| │ │ first and, among boosted, the │ | ||
| │ │ newest-arrived wins, so a │ |
| │ │ keys: sessionID │ │ (after response sent) │ │ | ||
| │ │ cap: 4096 default│ └─────────────────────────┘ │ | ||
| │ │ key: sessionID │ │ (after response sent) │ │ | ||
| │ │ value: completeAt│ └─────────────────────────┘ │ |
| 2. **Within the boosted group**, requests are ordered by their session's **completion time**: the session whose previous turn completed most recently is served first. Its prefix (KV) cache is the most likely to still be warm on the backend, so serving it first maximizes cache reuse. This ordering does not depend on when the follow-up request arrived, only on how recently the session's last turn finished; ties are broken FIFO by arrival time. | ||
| 3. **Among non-boosted requests**, ordering is FIFO (earliest arrival first). | ||
|
|
||
| Ordering by completion time (rather than a plain LIFO on arrival) keeps the queue fair: a boosted request cannot be starved by newer arrivals, because a session's completion time is fixed once its turn finishes and does not keep getting pushed back. Every boosted request also still leapfrogs the entire non-boosted tier, so a follow-up turn is always prioritized over unrelated traffic. |
Comment on lines
+528
to
+532
| // Count this boosted dequeue so the overtake bound advances for every | ||
| // still-waiting boosted request (see promoteStarvedBoostLocked). | ||
| if pq.sessionBoost && req.SessionBoost { | ||
| pq.boostAdmitCount++ | ||
| } |
Comment on lines
+166
to
+170
| │ │ head. Boosted requests rank │ | ||
| │ │ first and, among boosted, the │ | ||
| │ │ newest-arrived wins, so a │ | ||
| │ │ same-session follow-up that │ | ||
| │ │ arrived during the window is │ |
| 2. **Within the boosted group**, requests are ordered by their session's **completion time**: the session whose previous turn completed most recently is served first. Its prefix (KV) cache is the most likely to still be warm on the backend, so serving it first maximizes cache reuse. This ordering does not depend on when the follow-up request arrived, only on how recently the session's last turn finished; ties are broken FIFO by arrival time. | ||
| 3. **Among non-boosted requests**, ordering is FIFO (earliest arrival first). | ||
|
|
||
| Ordering by completion time (rather than a plain LIFO on arrival) keeps the queue fair: a boosted request cannot be starved by newer arrivals, because a session's completion time is fixed once its turn finishes and does not keep getting pushed back. Every boosted request also still leapfrogs the entire non-boosted tier, so a follow-up turn is always prioritized over unrelated traffic. |
Member
Author
|
@hzxuzhonghu Thanks for flagging the starvation problem. You're right: with plain newest-first ordering, every later-completing session gets a fresher timestamp and stays ahead of an older queued request (which keeps the timestamp captured at enqueue), so a steady stream of boosted traffic could push the older request back until it hit I've added bounded overtaking to fix this:
Test: Implementation note: to keep this cheap, the queue tracks overtakes with a single monotonic counter snapshotted per request (no per-enqueue heap scan) and promotes the forced request in place with Docs (proposal + user guide) updated with the new knob and rationale. All datastore tests pass under |
Member
Author
|
The CI failure is known issue fixed by: #1345 |
In session-boost mode, boosted requests were dequeued FIFO among themselves. Change the ordering so the most-recently-arrived boosted request wins (LIFO): a freshly boosted follow-up always jumps to the head of the queue so it reuses the still-warm prefix cache before earlier boosted requests. Non-boosted requests keep FIFO ordering, and boosted requests still outrank non-boosted ones. Also remove the grace-period 'head already boosted' fast path so the freed slot is held for the full grace window, allowing a still-newer same-session follow-up to arrive and take the head. Docs updated accordingly. Signed-off-by: YaoZengzeng <yaozengzeng@huawei.com>
Replace the LIFO-by-arrival ordering of boosted requests with an ordering based on each session's completion time: the session whose previous turn completed most recently is served first, because its prefix (KV) cache is the most likely to still be warm on the backend. Unlike plain LIFO, a session's completion time is fixed once its turn finishes, so a waiting boosted request cannot be starved by a stream of newer arrivals. Boosted requests still always outrank non-boosted ones, and non-boosted requests keep FIFO ordering. SessionTracker now stores the completion timestamp per session (LRU value) and exposes CompletionTime; PushRequest captures it into the new Request.SessionCompletedAt field used by Less. Signed-off-by: YaoZengzeng <yaozengzeng@huawei.com>
Address review feedback: SessionCompletedAt could be read as the whole session ending. The field is the completion time of the session's previous turn, so rename it to LastTurnCompletedAt for clarity. Signed-off-by: YaoZengzeng <yaozengzeng@huawei.com>
The completion-time ordering tests relied on consecutive MarkRequestCompleted calls producing strictly increasing time.Now() values. On a coarse clock or under fast execution two completions could share an instant, making the boosted ordering fall back to RequestTime and the expected order nondeterministic. Inject an overridable clock into SessionTracker (defaults to time.Now) and use it in the tests to assign strictly distinct completion timestamps, removing the flake without arbitrary sleeps. Signed-off-by: YaoZengzeng <yaozengzeng@huawei.com>
Signed-off-by: YaoZengzeng <yaozengzeng@huawei.com>
Completion-time ordering of boosted requests lets each later-completing session take a newer timestamp and stay ahead of an older queued request, which keeps the timestamp captured at enqueue. Under sustained boosted traffic the older request could be overtaken indefinitely and hit SESSION_BOOST_TIMEOUT instead of running. Bound the overtaking: each queued boosted request counts how many times a newer completion jumps ahead of it, and once that reaches SESSION_BOOST_MAX_OVERTAKES (default 16, configurable via env) it is marked aged and ranks ahead of the completion-time lane so it can no longer be pushed back. Add a test that injects a newer completion before every dequeue and asserts the victim is admitted within the bound, plus docs for the new knob. Signed-off-by: YaoZengzeng <yaozengzeng@huawei.com>
The first bounded-overtaking implementation scanned the whole heap on every boosted enqueue to bump a per-request overtake counter and rebuilt the heap with heap.Init, making enqueue O(n). Replace it with an O(log n) scheme: the queue keeps one monotonic counter of boosted dequeues, and each request snapshots it on arrival, so the difference is how many boosted requests were served ahead of it. Queued boosted requests are tracked in arrival order (lazy-skip list) so the longest waiter is found in amortized O(1); at dequeue, once its overtake count reaches SESSION_BOOST_MAX_OVERTAKES it is marked forced and promoted in place with heap.Fix. Enqueue is now a counter snapshot plus heap.Push (O(log n)); the dequeue check adds O(log n) only when a request must be promoted. Heap positions are tracked via heapIndex maintained in Swap/Push/Pop (and reset in the cancel drain) so heap.Fix targets the right slot. Signed-off-by: YaoZengzeng <yaozengzeng@huawei.com>
YaoZengzeng
force-pushed
the
feat/session-boost-newest-first
branch
from
d4f73d0 to
5d0b580
Compare
Comment on lines
+166
to
+170
| │ │ head. Boosted requests rank │ | ||
| │ │ first and, among boosted, the │ | ||
| │ │ newest-arrived wins, so a │ | ||
| │ │ same-session follow-up that │ | ||
| │ │ arrived during the window is │ |
| 2. **Within the boosted group**, requests are ordered by their session's **completion time**: the session whose previous turn completed most recently is served first. Its prefix (KV) cache is the most likely to still be warm on the backend, so serving it first maximizes cache reuse. This ordering does not depend on when the follow-up request arrived, only on how recently the session's last turn finished; ties are broken FIFO by arrival time. | ||
| 3. **Among non-boosted requests**, ordering is FIFO (earliest arrival first). | ||
|
|
||
| Ordering by completion time (rather than a plain LIFO on arrival) keeps the queue fair: a boosted request cannot be starved by newer arrivals, because a session's completion time is fixed once its turn finishes and does not keep getting pushed back. Every boosted request also still leapfrogs the entire non-boosted tier, so a follow-up turn is always prioritized over unrelated traffic. |
Comment on lines
+76
to
+80
| // SessionBoostMaxOvertakes bounds how many times a queued boosted request may be | ||
| // overtaken by newer same-priority boosted requests (sessions that completed | ||
| // more recently) before it is "aged": once a boosted request has been overtaken | ||
| // this many times it ranks ahead of the completion-time lane and can no longer | ||
| // be pushed back, so sustained boosted traffic cannot starve it into hitting |
Comment on lines
+187
to
+190
| // heap. boostEnqueueAdmits snapshots pq.boostAdmitCount at enqueue, so | ||
| // pq.boostAdmitCount-boostEnqueueAdmits is how many boosted requests have been | ||
| // served ahead of this one (i.e. how many times it has been overtaken). | ||
| // boostForced is set once that count reaches the bound; it makes the request |
Comment on lines
+434
to
+435
| pq.boostArrivals[0] = nil | ||
| pq.boostArrivals = pq.boostArrivals[1:] |
| | `ENABLE_SESSION_BOOST` | Enable the session-boost scheduling strategy | `false` | Mutually exclusive with `ENABLE_FAIRNESS_SCHEDULING` | | ||
| | `SESSION_BOOST_HEADER` | HTTP header used to identify conversation sessions | `X-Session-ID` | Must match what your clients send | | ||
| | `SESSION_BOOST_MAX_SESSIONS` | Max number of recently-completed sessions kept "warm" for boosting (LRU-bounded) | `4096` | When the cache is full, the least-recently-used session is evicted automatically. Size it by the number of concurrent conversations you want to keep boosted—no time-based tuning required | | ||
| | `SESSION_BOOST_MAX_OVERTAKES` | Max times a queued boosted request may be overtaken by newer completions | `16` | Bounds starvation: once a boosted request has been overtaken this many times it is "forced" and ranks ahead of newer completions, guaranteeing it runs before hitting `SESSION_BOOST_TIMEOUT`. Lower it to bound tail latency; raise it to favor prefix-cache reuse | |
| | `ENABLE_SESSION_BOOST` | `false` | Enable the session-boost scheduling strategy (mutually exclusive with `ENABLE_FAIRNESS_SCHEDULING`) | | ||
| | `SESSION_BOOST_HEADER` | `X-Session-ID` | HTTP header used to identify conversation sessions | | ||
| | `SESSION_BOOST_MAX_SESSIONS` | `4096` | Maximum number of recently-completed sessions kept warm for boosting. Bounds an LRU cache; the least-recently-used session is evicted when exceeded. Sized by session count, not time | | ||
| | `SESSION_BOOST_MAX_OVERTAKES` | `16` | Bounds how many times a queued boosted request may be overtaken by newer completions before it is "forced" and ranks ahead of the completion-time lane. Prevents sustained boosted traffic from starving an older request into hitting `SESSION_BOOST_TIMEOUT` | |
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
What type of PR is this?
/kind enhancement
What this PR does / why we need it:
In session-boost mode the per-model priority queue orders boosted requests
(follow-up turns of a recently-completed conversation) ahead of all non-boosted
traffic. This PR improves how boosted requests are ordered relative to each
other so that prefix-cache reuse is maximized without starving any request.
Boosted requests are now ordered by their session's completion time: the
session whose previous turn completed most recently is served first, because its
prefix (KV) cache is the most likely to still be warm on the backend. Ordering
depends only on how recently the session's last turn finished, not on when the
follow-up request arrived; ties are broken FIFO by arrival time.
Why completion time rather than a plain LIFO on arrival:
boosted request cannot be perpetually pushed back by a stream of newer
arrivals. Plain LIFO has no such bound and can starve an older boosted request.
degrades gracefully: once a session's cache is likely cold it also ages out of
the tracker's LRU, so it stops competing for the head.
requests keep FIFO ordering, so a follow-up turn is always prioritized over
unrelated traffic.
Changes:
session_boost_queue.go:SessionTrackerstores each session's completiontimestamp (LRU value) and exposes
CompletionTime.fairness_queue.go: newRequest.SessionCompletedAt;Lessorders boostedrequests by most-recent completion time (FIFO tiebreak);
PushRequestcaptures the completion time when boosting.
session_boost_queue_test.go: updated expectations and a newTestSessionBoostQueue_OrderedByCompletionTimethat distinguishescompletion-time ordering from both FIFO and LIFO.
session-boost.md(user guide) andsession-boost-strategy.md(proposal) updated to describe completion-time ordering.
Which issue(s) this PR fixes:
Fixes #
Special notes for your reviewer:
Unit tests in
pkg/kthena-router/datastorepass locally (go test ./pkg/kthena-router/datastore/...).This PR was prepared with AI assistance; I have reviewed and understand every change.
Does this PR introduce a user-facing change?: